#include "U2CodecFreeImage.h"

#include "U2LogManager.h"
#include "U2StringUtil.h"
#include "U2Exception.h"
#include "FreeImage/FreeImage.h"



// freeimage 3.9.1~3.11.0 interoperability fix
#ifndef FREEIMAGE_COLORORDER
    // we have freeimage 3.9.1, define these symbols in such way as 3.9.1 really work (do not use 3.11.0 definition, as color order was changed between these two versions on Apple systems)
    #define FREEIMAGE_COLORORDER_BGR	0
    #define FREEIMAGE_COLORORDER_RGB	1
    #if defined(FREEIMAGE_BIGENDIAN)
        #define FREEIMAGE_COLORORDER    FREEIMAGE_COLORORDER_RGB
    #else
        #define FREEIMAGE_COLORORDER    FREEIMAGE_COLORORDER_BGR
    #endif
#endif



U2EG_NAMESPACE_USING


U2CodecFreeImage::RegisteredCodecList U2CodecFreeImage::s_codecList;
//---------------------------------------------------------------------
void FreeImageErrorHandler(FREE_IMAGE_FORMAT fif, const char *message) 
{
    // Callback method as required by FreeImage to report problems
    const char* typeName = FreeImage_GetFormatFromFIF(fif);
    if (typeName)
    {
        U2LogManager::getSingleton().stream() 
            << "FreeImage error: '" << message << "' when loading format "
            << typeName;
    }
    else
    {
        U2LogManager::getSingleton().stream() 
            << "FreeImage error: '" << message << "'";
    }
}
//---------------------------------------------------------------------
void U2CodecFreeImage::startup(void)
{
    FreeImage_Initialise(false);

    U2LogManager::getSingleton().logMessage(
        LML_NORMAL,
        "FreeImage version: " + U2String(FreeImage_GetVersion()));
    U2LogManager::getSingleton().logMessage(
        LML_NORMAL,
        FreeImage_GetCopyrightMessage());

    // Register codecs
    U2StringUtil::U2StrStreamType strExt;
    strExt << "Supported formats: ";
    bool first = true;
    for (int i = 0; i < FreeImage_GetFIFCount(); ++i)
    {

        // Skip DDS codec since FreeImage does not have the option 
        // to keep DXT data compressed, we'll use our own codec
        if ((FREE_IMAGE_FORMAT)i == FIF_DDS)
            continue;

        U2String exts(FreeImage_GetFIFExtensionList((FREE_IMAGE_FORMAT)i));
        if (!first)
        {
            strExt << ",";
        }
        first = false;
        strExt << exts;

        // Pull off individual formats (separated by comma by FI)
        U2StringVector extsVector = U2StringUtil::split(exts, ",");
        for (U2StringVector::iterator v = extsVector.begin(); v != extsVector.end(); ++v)
        {
            U2CodecImage* codec = U2_NEW U2CodecFreeImage(*v, i);
            s_codecList.push_back(codec);
            U2Codec::registerCodec(codec);
        }
    }
    U2LogManager::getSingleton().logMessage(
        LML_NORMAL,
        strExt.str());

    // Set error handler
    FreeImage_SetOutputMessage(FreeImageErrorHandler);
}
//---------------------------------------------------------------------
void U2CodecFreeImage::shutdown(void)
{
    FreeImage_DeInitialise();

    for (RegisteredCodecList::iterator i = s_codecList.begin();
        i != s_codecList.end(); ++i)
    {
        U2Codec::unRegisterCodec(*i);
        U2_DELETE *i;
    }
    s_codecList.clear();

}
//---------------------------------------------------------------------
U2CodecFreeImage::U2CodecFreeImage(const U2String &type, unsigned int fiType):
m_szType(type),
m_unFreeImageType(fiType)
{ 
}
//---------------------------------------------------------------------
FIBITMAP* U2CodecFreeImage::encode(U2DataStreamMemoryPtr& input, CodecDataPtr& pData) const
{
    FIBITMAP* ret = 0;

    ImageData* pImgData = static_cast< ImageData * >( pData.getPointer() );
    U2PixelBox src(pImgData->width, pImgData->height, pImgData->depth, pImgData->format, input->getPtr());

    // The required format, which will adjust to the format
    // actually supported by FreeImage.
    PixelFormat requiredFormat = pImgData->format;

    // determine the settings
    FREE_IMAGE_TYPE imageType;
    PixelFormat determiningFormat = pImgData->format;

    switch(determiningFormat)
    {
    case PF_R5G6B5:
    case PF_B5G6R5:
    case PF_R8G8B8:
    case PF_B8G8R8:
    case PF_A8R8G8B8:
    case PF_X8R8G8B8:
    case PF_A8B8G8R8:
    case PF_X8B8G8R8:
    case PF_B8G8R8A8:
    case PF_R8G8B8A8:
    case PF_A4L4:
    case PF_BYTE_LA:
    case PF_R3G3B2:
    case PF_A4R4G4B4:
    case PF_A1R5G5B5:
    case PF_A2R10G10B10:
    case PF_A2B10G10R10:
        // I'd like to be able to use r/g/b masks to get FreeImage to load the data
        // in it's existing format, but that doesn't work, FreeImage needs to have
        // data in RGB[A] (big endian) and BGR[A] (little endian), always.
        if (U2PixelUtil::hasAlpha(determiningFormat))
        {
#if FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_RGB
            requiredFormat = PF_BYTE_RGBA;
#else
            requiredFormat = PF_BYTE_BGRA;
#endif
        }
        else
        {
#if FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_RGB
            requiredFormat = PF_BYTE_RGB;
#else
            requiredFormat = PF_BYTE_BGR;
#endif
        }
        // fall through
    case PF_L8:
    case PF_A8:
        imageType = FIT_BITMAP;
        break;

    case PF_L16:
        imageType = FIT_UINT16;
        break;

    case PF_SHORT_GR:
        requiredFormat = PF_SHORT_RGB;
        // fall through
    case PF_SHORT_RGB:
        imageType = FIT_RGB16;
        break;

    case PF_SHORT_RGBA:
        imageType = FIT_RGBA16;
        break;

    case PF_FLOAT16_R:
        requiredFormat = PF_FLOAT32_R;
        // fall through
    case PF_FLOAT32_R:
        imageType = FIT_FLOAT;
        break;

    case PF_FLOAT16_GR:
    case PF_FLOAT16_RGB:
    case PF_FLOAT32_GR:
        requiredFormat = PF_FLOAT32_RGB;
        // fall through
    case PF_FLOAT32_RGB:
        imageType = FIT_RGBF;
        break;

    case PF_FLOAT16_RGBA:
        requiredFormat = PF_FLOAT32_RGBA;
        // fall through
    case PF_FLOAT32_RGBA:
        imageType = FIT_RGBAF;
        break;

    default:
        U2_EXCEPT(U2Exception::ERR_ITEM_NOT_FOUND, "Invalid image format", "U2CodecFreeImage::encode");
    };

    // Check support for this image type & bit depth
    if (!FreeImage_FIFSupportsExportType((FREE_IMAGE_FORMAT)m_unFreeImageType, imageType) ||
        !FreeImage_FIFSupportsExportBPP((FREE_IMAGE_FORMAT)m_unFreeImageType, (int)U2PixelUtil::getNumElemBits(requiredFormat)))
    {
        // Ok, need to allocate a fallback
        // Only deal with RGBA -> RGB for now
        switch (requiredFormat)
        {
        case PF_BYTE_RGBA:
            requiredFormat = PF_BYTE_RGB;
            break;
        case PF_BYTE_BGRA:
            requiredFormat = PF_BYTE_BGR;
            break;
        default:
        	;
        };

    }

    bool conversionRequired = false;

    unsigned char* srcData = input->getPtr();

    // Check BPP
    unsigned bpp = static_cast<unsigned>(U2PixelUtil::getNumElemBits(requiredFormat));
    if (!FreeImage_FIFSupportsExportBPP((FREE_IMAGE_FORMAT)m_unFreeImageType, (int)bpp))
    {
        if (bpp == 32 && U2PixelUtil::hasAlpha(pImgData->format) && FreeImage_FIFSupportsExportBPP((FREE_IMAGE_FORMAT)m_unFreeImageType, 24))
        {
            // drop to 24 bit (lose alpha)
#if FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_RGB
            requiredFormat = PF_BYTE_RGB;
#else
            requiredFormat = PF_BYTE_BGR;
#endif
            bpp = 24;
        }
        else if (bpp == 128 && U2PixelUtil::hasAlpha(pImgData->format) && FreeImage_FIFSupportsExportBPP((FREE_IMAGE_FORMAT)m_unFreeImageType, 96))
        {
            // drop to 96-bit floating point
            requiredFormat = PF_FLOAT32_RGB;
        }
    }

    U2PixelBox convBox(pImgData->width, pImgData->height, 1, requiredFormat);
    if (requiredFormat != pImgData->format)
    {
        conversionRequired = true;
        // Allocate memory
        convBox.data = U2_ALLOC_T(u2uchar, convBox.getConsecutiveSize(), MEMCATEGORY_GENERAL);
        // perform conversion and reassign source
        U2PixelBox src(pImgData->width, pImgData->height, 1, pImgData->format, input->getPtr());
        U2PixelUtil::bulkPixelConversion(src, convBox);
        srcData = static_cast<unsigned char*>(convBox.data);

    }


    ret = FreeImage_AllocateT(
        imageType, 
        static_cast<int>(pImgData->width), 
        static_cast<int>(pImgData->height), 
        bpp);

    if (!ret)
    {
        if (conversionRequired)
            U2_FREE(convBox.data, MEMCATEGORY_GENERAL);

        U2_EXCEPT(U2Exception::ERR_INVALIDPARAMS, 
            "FreeImage_AllocateT failed - possibly out of memory. ", 
            __FUNCTION__);
    }

    if (requiredFormat == PF_L8 || requiredFormat == PF_A8)
    {
        // Must explicitly tell FreeImage that this is greyscale by setting
        // a "grey" palette (otherwise it will save as a normal RGB
        // palettized image).
        FIBITMAP *tmp = FreeImage_ConvertToGreyscale(ret);
        FreeImage_Unload(ret);
        ret = tmp;
    }

    size_t dstPitch = FreeImage_GetPitch(ret);
    size_t srcPitch = pImgData->width * U2PixelUtil::getNumElemBytes(requiredFormat);


    // Copy data, invert scanlines and respect FreeImage pitch
    u2uchar* pSrc;
    u2uchar* pDst = FreeImage_GetBits(ret);
    for (size_t y = 0; y < pImgData->height; ++y)
    {
        pSrc = srcData + (pImgData->height - y - 1) * srcPitch;
        memcpy(pDst, pSrc, srcPitch);
        pDst += dstPitch;
    }

    if (conversionRequired)
    {
        // delete temporary conversion area
        U2_FREE(convBox.data, MEMCATEGORY_GENERAL);
    }

    return ret;
}
//---------------------------------------------------------------------
U2DataStreamPtr U2CodecFreeImage::code(U2DataStreamMemoryPtr& input, U2Codec::CodecDataPtr& pData) const
{        
    FIBITMAP* fiBitmap = encode(input, pData);

    // open memory chunk allocated by FreeImage
    FIMEMORY* mem = FreeImage_OpenMemory();
    // write data into memory
    FreeImage_SaveToMemory((FREE_IMAGE_FORMAT)m_unFreeImageType, fiBitmap, mem);
    // Grab data information
    BYTE* data;
    DWORD size;
    FreeImage_AcquireMemory(mem, &data, &size);
    // Copy data into our own buffer
    // Because we're asking U2DataStreamMemory to free this, must create in a compatible way
    BYTE* ourData = U2_ALLOC_T(BYTE, size, MEMCATEGORY_GENERAL);
    memcpy(ourData, data, size);
    // Wrap data in stream, tell it to free on close 
    U2DataStreamPtr outstream(U2_NEW U2DataStreamMemory(ourData, size, true));
    // Now free FreeImage memory buffers
    FreeImage_CloseMemory(mem);
    // Unload bitmap
    FreeImage_Unload(fiBitmap);

    return outstream;


}
//---------------------------------------------------------------------
void U2CodecFreeImage::codeToFile(U2DataStreamMemoryPtr& input, 
                                const U2String& outFileName, U2Codec::CodecDataPtr& pData) const
{
    FIBITMAP* fiBitmap = encode(input, pData);

    FreeImage_Save((FREE_IMAGE_FORMAT)m_unFreeImageType, fiBitmap, outFileName.c_str());
    FreeImage_Unload(fiBitmap);


}
//---------------------------------------------------------------------
U2Codec::DecodeResult U2CodecFreeImage::decode(U2DataStreamPtr& input) const
{
    // Buffer stream into memory (TODO: override IO functions instead?)
    U2DataStreamMemory memStream(input, true);

    FIMEMORY* fiMem = 
        FreeImage_OpenMemory(memStream.getPtr(), static_cast<DWORD>(memStream.size()));

    FIBITMAP* fiBitmap = FreeImage_LoadFromMemory(
        (FREE_IMAGE_FORMAT)m_unFreeImageType, fiMem);
    if (!fiBitmap)
    {
        U2_EXCEPT(U2Exception::ERR_INTERNAL_ERROR, 
            "Error decoding image", 
            "U2CodecFreeImage::decode");
    }


    ImageData* imgData = U2_NEW ImageData();
    U2DataStreamMemoryPtr output;

    imgData->depth = 1; // only 2D formats handled by this codec
    imgData->width = FreeImage_GetWidth(fiBitmap);
    imgData->height = FreeImage_GetHeight(fiBitmap);
    imgData->num_mipmaps = 0; // no mipmaps in non-DDS 
    imgData->flags = 0;

    // Must derive format first, this may perform conversions

    FREE_IMAGE_TYPE imageType = FreeImage_GetImageType(fiBitmap);
    FREE_IMAGE_COLOR_TYPE colourType = FreeImage_GetColorType(fiBitmap);
    unsigned bpp = FreeImage_GetBPP(fiBitmap);

    switch(imageType)
    {
    case FIT_UNKNOWN:
    case FIT_COMPLEX:
    case FIT_UINT32:
    case FIT_INT32:
    case FIT_DOUBLE:
    default:
        U2_EXCEPT(U2Exception::ERR_ITEM_NOT_FOUND, 
            "Unknown or unsupported image format", 
            "U2CodecFreeImage::decode");

        break;
    case FIT_BITMAP:
        // Standard image type
        // Perform any colour conversions for greyscale
        if (colourType == FIC_MINISWHITE || colourType == FIC_MINISBLACK)
        {
            FIBITMAP* newBitmap = FreeImage_ConvertToGreyscale(fiBitmap);
            // free old bitmap and replace
            FreeImage_Unload(fiBitmap);
            fiBitmap = newBitmap;
            // get new formats
            bpp = FreeImage_GetBPP(fiBitmap);
            colourType = FreeImage_GetColorType(fiBitmap);
        }
        // Perform any colour conversions for RGB
        else if (bpp < 8 || colourType == FIC_PALETTE || colourType == FIC_CMYK)
        {
            FIBITMAP* newBitmap = FreeImage_ConvertTo24Bits(fiBitmap);
            // free old bitmap and replace
            FreeImage_Unload(fiBitmap);
            fiBitmap = newBitmap;
            // get new formats
            bpp = FreeImage_GetBPP(fiBitmap);
            colourType = FreeImage_GetColorType(fiBitmap);
        }

        // by this stage, 8-bit is greyscale, 16/24/32 bit are RGB[A]
        switch(bpp)
        {
        case 8:
            imgData->format = PF_L8;
            break;
        case 16:
            // Determine 555 or 565 from green mask
            // cannot be 16-bit greyscale since that's FIT_UINT16
            if(FreeImage_GetGreenMask(fiBitmap) == FI16_565_GREEN_MASK)
            {
                imgData->format = PF_R5G6B5;
            }
            else
            {
                // FreeImage doesn't support 4444 format so must be 1555
                imgData->format = PF_A1R5G5B5;
            }
            break;
        case 24:
            // FreeImage differs per platform
            //     PF_BYTE_BGR[A] for little endian (== PF_ARGB native)
            //     PF_BYTE_RGB[A] for big endian (== PF_RGBA native)
#if FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_RGB
            imgData->format = PF_BYTE_RGB;
#else
            imgData->format = PF_BYTE_BGR;
#endif
            break;
        case 32:
#if FREEIMAGE_COLORORDER == FREEIMAGE_COLORORDER_RGB
            imgData->format = PF_BYTE_RGBA;
#else
            imgData->format = PF_BYTE_BGRA;
#endif
            break;


        };
        break;
    case FIT_UINT16:
    case FIT_INT16:
        // 16-bit greyscale
        imgData->format = PF_L16;
        break;
    case FIT_FLOAT:
        // Single-component floating point data
        imgData->format = PF_FLOAT32_R;
        break;
    case FIT_RGB16:
        imgData->format = PF_SHORT_RGB;
        break;
    case FIT_RGBA16:
        imgData->format = PF_SHORT_RGBA;
        break;
    case FIT_RGBF:
        imgData->format = PF_FLOAT32_RGB;
        break;
    case FIT_RGBAF:
        imgData->format = PF_FLOAT32_RGBA;
        break;


    };

    unsigned char* srcData = FreeImage_GetBits(fiBitmap);
    unsigned srcPitch = FreeImage_GetPitch(fiBitmap);

    // Final data - invert image and trim pitch at the same time
    size_t dstPitch = imgData->width * U2PixelUtil::getNumElemBytes(imgData->format);
    imgData->size = dstPitch * imgData->height;
    // Bind output buffer
    output.bind(U2_NEW U2DataStreamMemory(imgData->size));

    u2uchar* pSrc;
    u2uchar* pDst = output->getPtr();
    for (size_t y = 0; y < imgData->height; ++y)
    {
        pSrc = srcData + (imgData->height - y - 1) * srcPitch;
        memcpy(pDst, pSrc, dstPitch);
        pDst += dstPitch;
    }


    FreeImage_Unload(fiBitmap);
    FreeImage_CloseMemory(fiMem);


    DecodeResult ret;
    ret.first = output;
    ret.second = CodecDataPtr(imgData);
    return ret;

}
//---------------------------------------------------------------------    
U2String U2CodecFreeImage::getType() const 
{
    return m_szType;
}
//---------------------------------------------------------------------
U2String U2CodecFreeImage::magicNumberToFileExt(const char *magicNumberPtr, size_t maxbytes) const
{
    FIMEMORY* fiMem = 
        FreeImage_OpenMemory((BYTE*)magicNumberPtr, static_cast<DWORD>(maxbytes));

    FREE_IMAGE_FORMAT fif = FreeImage_GetFileTypeFromMemory(fiMem, (int)maxbytes);
    FreeImage_CloseMemory(fiMem);

    if (fif != FIF_UNKNOWN)
    {
        U2String ext(FreeImage_GetFormatFromFIF(fif));
        U2StringUtil::toLowerCase(ext);
        return ext;
    }
    else
    {
        return U2StringUtil::BLANK;
    }
}
